package com.hy.aicodegeneration.pool;

import jakarta.annotation.PreDestroy;
import lombok.extern.slf4j.Slf4j;
import org.openqa.selenium.chrome.ChromeDriver;
import org.openqa.selenium.chrome.ChromeOptions;
import org.springframework.stereotype.Component;

import java.util.Queue;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * WebDriver 连接池实现
 */
@Slf4j
@Component
public class WebDriverPool {

    private final Queue<PooledWebDriver> pool = new ConcurrentLinkedQueue<>();
    private final AtomicInteger activeCount = new AtomicInteger(0);
    private final int maxPoolSize = 5; // 最大连接数

    public WebDriverPool() {
        initializePool();
    }

    /**
     * 初始化连接池
     */
    private void initializePool() {
        for (int i = 0; i < maxPoolSize; i++) {
            PooledWebDriver driver = new PooledWebDriver(createChromeDriver(), this);
            pool.offer(driver);
        }
        log.info("WebDriver 连接池初始化完成，初始容量: {}", maxPoolSize);
    }

    /**
     * 获取空闲的 WebDriver 实例
     *
     * @return PooledWebDriver 实例
     */
    public PooledWebDriver borrowDriver() {
        PooledWebDriver driver = pool.poll();
        
        // 如果池中没有可用实例，则创建新的（不超过最大限制）
        if (driver == null && activeCount.get() < maxPoolSize * 2) {
            driver = new PooledWebDriver(createChromeDriver(), this);
            log.debug("创建新的 WebDriver 实例，当前活跃数: {}", activeCount.incrementAndGet());
        } else if (driver != null) {
            activeCount.incrementAndGet();
            log.debug("从池中获取 WebDriver 实例，当前活跃数: {}", activeCount.get());
        }

        return driver;
    }

    /**
     * 归还 WebDriver 实例到池中
     *
     * @param driver 要归还的 WebDriver 实例
     */
    public void returnDriver(PooledWebDriver driver) {
        if (driver != null) {
            try {
                // 清理当前会话状态
                driver.getWebDriver().manage().deleteAllCookies();
                driver.getWebDriver().navigate().to("about:blank");
                log.debug("清理 WebDriver 状态完成");
            } catch (Exception e) {
                log.warn("清理 WebDriver 状态失败", e);
            }

            // 如果池未满则归还，否则关闭
            if (pool.size() < maxPoolSize) {
                pool.offer(driver);
                log.debug("WebDriver 实例归还到池中");
            } else {
                try {
                    driver.getWebDriver().quit();
                    log.debug("WebDriver 实例已关闭（池已满）");
                } catch (Exception e) {
                    log.warn("关闭 WebDriver 实例失败", e);
                }
            }
        }
        
        activeCount.decrementAndGet();
        log.debug("归还 WebDriver 实例，当前活跃数: {}", activeCount.get());
    }

    /**
     * 创建 ChromeDriver 实例
     *
     * @return ChromeDriver 实例
     */
    private ChromeDriver createChromeDriver() {
        ChromeOptions options = new ChromeOptions();
        options.addArguments("--headless");
        options.addArguments("--no-sandbox");
        options.addArguments("--disable-dev-shm-usage");
        options.addArguments("--disable-gpu");
        options.addArguments("--window-size=1920,1080");
        options.addArguments("--disable-extensions");
        options.addArguments("--disable-plugins");
        options.addArguments("--disable-images"); // 可选：禁用图片加载提升速度
        options.addArguments("--disable-javascript"); // 可选：禁用 JS 提升速度（如果不需要交互）

        return new ChromeDriver(options);
    }

    /**
     * 获取当前活跃的 WebDriver 数量
     *
     * @return 活跃数量
     */
    public int getActiveCount() {
        return activeCount.get();
    }

    /**
     * 获取池中可用实例数量
     *
     * @return 可用数量
     */
    public int getAvailableCount() {
        return pool.size();
    }

    /**
     * 获取池的最大容量
     *
     * @return 最大容量
     */
    public int getMaxPoolSize() {
        return maxPoolSize;
    }

    /**
     * 关闭所有 WebDriver 实例（应用关闭时调用）
     */
    @PreDestroy
    public void shutdown() {
        log.info("开始关闭 WebDriver 连接池...");
        while (!pool.isEmpty()) {
            PooledWebDriver driver = pool.poll();
            if (driver != null) {
                try {
                    driver.getWebDriver().quit();
                } catch (Exception e) {
                    log.warn("关闭 WebDriver 实例失败", e);
                }
            }
        }
        log.info("WebDriver 连接池已关闭");
    }
}
